import pandas as pd
import xlrd
import openpyxl as ol
import os.path


class NotFoundHdrFlag(Exception):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return self.msg


class HdrTypeError(Exception):
    def __init__(self, msg):
        self._msg = msg

    def __str__(self):
        return self._msg


class ExcelMultiLineHdrDetect():
    """
    Excel 表中，表头跨越多行情况的处理。可以根据起始位置唯一标识和结束位置唯一标识判定表头中列名区域的位置。
    """

    def __init__(self, path, sheetName=None, n_row=10):
        """
        :param path:文件路径
        :param sheetName: Excel的Sheet表名。
        :param n_row: 表头范围检测行数，只在指定的前n_row行中进行检测。
        """
        self._path = path
        self._sheetName = sheetName
        self._n_row = n_row
        # 读取Excel文件中sheetName对应表的合并单元格的位置
        self._merged_list = self._read_excel_mergeds(sheetName)
        # 使用Pandas读取SheetName表中的头部区域
        self._hdr_data = self._readData()

    def getModifyData(self, start_flag="序号", end_flag="备注"):
        columns_data = self.getColumns(start_flag, end_flag)
        modify_map = self.getModifyMap(start_flag, end_flag)
        return columns_data.loc[:, modify_map]

    def getModifyMap(self, start_flag="序号", end_flag="备注"):
        """
        获取每列标题中的列名是否需要更详细修改的映射表。
        :param start_flag: 用于定位文件头部内容的唯一性标识。
        :param end_flag: Sheet表中头部区域标题的末尾位置处的标识文字。
        :return: 一个与列名个数相同的True及False列表。
        """
        start_index, start_col, end_index, end_col = self.getColumnsPos(start_flag, end_flag)

        # 查找所有在此范围内的合并单元格，如果此合并单元格跨列了，则需要进行详细处理，如果仅仅是跨行而不跨列，则不需要详细处理
        columns_map = [False] * (end_col - start_col + 1)
        for X_s, Y_s, X_e, Y_e in self._merged_list:
            if start_col <= X_s and X_e <= end_col and start_index <= Y_s and Y_e <= end_index:
                if X_s != X_e:
                    for col in range(X_s, X_e + 1):
                        columns_map[col - start_col] = True

        return columns_map

    def _readData(self):
        """使用Pandas读取Sheet表的前n_row行进行分析"""
        hdr_data = pd.read_excel(self._path, sheet_name=self._sheetName, header=None, nrows=self._n_row)
        return hdr_data

    def getColumns(self, start_flag="序号", end_flag="备注"):
        """
        根据开始标识和结束标识确定跨多行及多列的Columns区域的位置。
        :param start_flag: 用于定位文件头部内容的唯一性标识。
        :param end_flag: Sheet表中头部区域标题的末尾位置处的标识文字。
        :return: df
        """
        start_index, start_col, end_index, end_col = self.getColumnsPos(start_flag, end_flag)
        return self._hdr_data.iloc[start_index:end_index + 1, start_col:end_col + 1]

    def getColumnsPos(self, start_flag="序号", end_flag="备注"):
        """
        根据指定的表头区域，column开始标识和结束标识确定跨多行column及多列的Columns区域的位置。
        :param start_flag: 用于定位文件头部内容的唯一性标识。
        :param end_flag: Sheet表中头部区域标题的末尾位置处的标识文字。
        :return: row_start,column_start,row_end,column_end
        """
        start_index, start_col = self.detectHdr(start_flag)
        end_index, end_col = self.detectTail(end_flag)
        return start_index, start_col, end_index, end_col

    def getHdrRowNum(self, uniqe_flag='序号'):
        """
        根据指定的标头区域唯一识别符号获取表头区域起始位置，并评估表头跨行数量。
        :param uniqe_flag: 用于定位文件头部内容的唯一性标识。
        :return: 如果返回为0，表示未找到指定标识文字（一般不会运行到，因为如果找不到会产生HdrTypeError异常；其它值表示表头可能跨越的行数。
        """
        pos = self.detectHdr(uniqe_flag)
        Y_hdr, X_hdr = pos
        span_row = 0
        for X_s, Y_s, X_e, Y_e in self._merged_list:
            if X_s <= X_hdr <= X_e and Y_s <= Y_hdr <= Y_e:
                span_row = Y_e - Y_s
        return span_row + 1

    def getRowNum(self, uniqe_flag='备注'):
        """
        根据Sheet表表头尾部标识获取头部区域的准确跨越行数，比getHdrRowNum更加精确。
        :param uniqe_flag:Sheet表中头部区域标题的末尾位置处的标识文字。
        """
        pos = self.detectTail(uniqe_flag)
        Y_tail, X_tail = pos
        span_row = 0
        for X_s, Y_s, X_e, Y_e in self._merged_list:
            if X_s <= X_tail <= X_e and Y_s <= Y_tail <= Y_e:
                span_row = Y_e - Y_s
        return span_row + 1

    def detectTail(self, uniqe_flag='备注'):
        """
        检测Excel的Sheet表中头部区域中指定标识所定位的尾部位置。
        :param uniqe_flag: Sheet表中头部区域标题的末尾位置处的标识文字。
        :return: 标志位置（row,col)，否则为（None，None）。注意，如果标识文字处理跨越几行的合并单元格中，则取此单元格右下角的位置。
        """
        # 先使用detectHdr方法获取标志的准确位置
        try:
            index_i, columns_i = self.detectHdr(uniqe_flag)
        except NotFoundHdrFlag:
            raise NotFoundHdrFlag("文件尾部未发现定位标识：[{}].".format(uniqe_flag))
        # 对于尾部位置需要进行处理，如果尾部标志位置处于合并单元格中，则要取合并单元格的右下角位置
        for X_s, Y_s, X_e, Y_e in self._merged_list:
            if Y_s <= index_i <= Y_e and X_s <= columns_i <= X_e:
                return Y_e, X_e
        return index_i, columns_i

    def detectHdr(self, uniqe_flag='序号'):
        """
        检测Excel头部标志。
        :param uniqe_flag:用于定位文件头部内容的唯一性标识。
        :return: 标志位置（row,col）,否则为（None，None）
        """
        pos_init = self._hdr_data[self._hdr_data == uniqe_flag]
        pos = pos_init.dropna(how='all')
        if pos.empty:
            raise NotFoundHdrFlag("文件【{}】中【{}】表头部未发现定位标识：[{}].".format(self._path,self._sheetName,uniqe_flag))
        else:
            index_str = pos.index.values[0]
            index_i = pos_init.index.tolist().index(index_str)
            columns_notna = pos.dropna(axis=1)
            columns_str = columns_notna.columns.values[0]
            columns_i = pos.columns.tolist().index(columns_str)
        return index_i, columns_i

    def _read_excel_mergeds(self, sheetName):
        """
        分别使用xlrd和openpyxl读取xls及xlsx格式的文件，这里主要用于获取表中的合并项，否则可以使用Pandas一次读取。
        经过测试，使用openpyxl读取数据文件时，可能会一次性读取非常多的无用列，故次方法中只读取合并单元格的信息。
        :param sheetName: 要读取的Excel表名。
        :return:
        """
        self._merged_list = []
        if os.path.splitext(self._path)[1] == '.xlsx':
            # 使用openpyxl读取xlsx文件
            self._wb = ol.load_workbook(self._path)
            # 获取此文件中所有工作表的名称
            self._sheet_names = self._wb.sheetnames
            # 判断如果没有指定表名称，则默认读取第一个Sheet表
            if sheetName is None:
                self._sheetName = self._sheet_names[0]
            # 如果传入的值为纯数字，则将其视为要读取的表的索引
            if isinstance(sheetName, int):
                self._sheetName = self._sheet_names[sheetName]
            # 获取Sheet表数据
            self._sheet = self._wb[self._sheetName]

            # 获取此Sheet表中，所有的合并单元格
            for range in self._sheet.merged_cells:
                self._merged_list.append(range.bounds)

            self._wb.close()
            # 对合并单元格进行转化，以便与使用xlrd读取方式获取的一致
            self._merged_list = [(X_s - 1, Y_s - 1, X_e - 1, Y_e - 1) for (X_s, Y_s, X_e, Y_e) in self._merged_list]
        elif os.path.splitext(self._path)[1] == '.xls':
            # 使用xlrd读取xlsx文件
            self._wb = xlrd.open_workbook(self._path, formatting_info=True)
            # 获取此文件中所有工作表的名称
            self._sheet_names = self._wb.sheet_names()
            # 判断如果没有指定表名称，则默认读取第一个Sheet表
            if sheetName is None:
                self._sheetName = self._sheet_names[0]
            # 如果传入的值为纯数字，则将其视为要读取的表的索引
            if isinstance(sheetName, int):
                self._sheetName = self._sheet_names[sheetName]
            # 获取Sheet表数据
            self._sheet = self._wb[self._sheetName]
            # 对合并单元格进行转化，以便与使用openpyxl读取方式获取的一致
            self._merged_list = [(X_s, Y_s, X_e - 1, Y_e - 1) for (Y_s, Y_e, X_s, X_e) in self._sheet.merged_cells]
        return self._merged_list

if __name__ == '__main__':
    path = "../原始数据/服务价格/21价格8.11.xlsx"
    ehd = ExcelMultiLineHdrDetect(path, sheetName=None)
    print("起始位置：", ehd.detectHdr("国家项目编码"))
    print("结束位置：", ehd.detectTail("支付类别"))
    print("表头部分列名区域：",ehd.getColumnsPos("国家项目编码","支付类别"))